﻿using System.Collections.Generic;
using HAPI;
using HuginApiAddonsCS.Policy;
using HuginApiAddonsCS.Policy.PolicyGenerator;
using HuginApiAddonsCS.Simulator;

namespace HuginApiAddonsCS.Extensions
{
    public static class DomainExtensions
    {
        /// <summary>
        /// Gets the first time step value from the domain
        /// </summary>
        /// <param name="domain"></param>
        /// <returns>value of first time step</returns>
        public static int GetFirstTimeStep(this Domain domain)
        {
            int min = -1;

            foreach (Node node in domain.GetNodes())
            {
                int t = node.GetNodeTimeStep();
                if (t < min || min == -1)
                {
                    min = t;
                }
            }

            return min;
        }

        /// <summary>
        /// Gets the value of the final time step
        /// </summary>
        /// <param name="domain"></param>
        /// <returns>the value of the final time step</returns>
        public static int GetLastTimeStep(this Domain domain)
        {
            int max = -1;

            foreach (Node node in domain.GetNodes())
            {
                int t = node.GetNodeTimeStep();
                if (t > max)
                {
                    max = t;
                }
            }

            return max;
        }

        /// <summary>
        /// Gets a list of all the prefixes in the domain
        /// </summary>
        /// <param name="domain"></param>
        /// <returns>list of prefixes</returns>
        public static List<string> GetPrefixes(this Domain domain)
        {
            Dictionary<string, string> domainPrefixes = new Dictionary<string, string>();

            foreach (Node node in domain.GetNodes())
            {
                string nodeName = node.GetNodePrefix();
                if (!domainPrefixes.ContainsKey(nodeName))
                {
                    domainPrefixes.Add(nodeName, nodeName);
                }
            }

            return new List<string>(domainPrefixes.Keys);
        }

        /// <summary>
        /// Gets the next node in the domain based on the current node
        /// </summary>
        /// <param name="domain">Domain extended</param>
        /// <param name="node">Node</param>
        /// <returns>The next node</returns>
        public static Node GetNextStepNode(this Domain domain, Node node)
        {
            return domain.GetNodeByName(node.GetNodePrefix() + (node.GetNodeTimeStep() + 1));
        }

        /// <summary>
        /// Gets the policy tree associated with theis domain
        /// </summary>
        /// <param name="domain">domain</param>
        /// <param name="ajPrefix">action prefix</param>
        /// <param name="ojPrefix">observation prefix</param>
        /// <returns>generated policy tree</returns>
        public static PolicyNode GetPolicyTree(this Domain domain, string ajPrefix, string ojPrefix)
        {
            PolicyGenerator generator = new PolicyGenerator();
            return generator.GeneratePolicyTree(domain, ajPrefix, ojPrefix);
        }

        /// <summary>
        /// Executes a single simulation on the domain
        /// </summary>
        /// <param name="domain">domain</param>
        /// <param name="ajPrefix">action prefix</param>
        /// <param name="ojPrefix">observation prefix</param>
        /// <param name="sPrefix">state prefix</param>
        /// <param name="uPrefix">utility prefix</param>
        /// <param name="model">model name</param>
        /// <returns>list of results</returns>
        public static List<SimulationResult> GetSimulationResultDid(this Domain domain, string ajPrefix, string ojPrefix,
            string sPrefix, string uPrefix, string model)
        {
            List<SimulationResult> results = new List<SimulationResult>();

            List<uint> stateHistory = new List<uint>();

            for (int i = GetFirstTimeStep(domain); i <= GetLastTimeStep(domain); i++)
            {
                SimulationResult result = new SimulationResult();
                result.model = model;
                result.timeStep = i;

                DiscreteChanceNode stateNode = (DiscreteChanceNode) domain.GetNodeByName(sPrefix + i);
                uint stateIndex = stateNode.SetRandomState();
                stateHistory.Add(stateIndex);
                result.state = stateNode.GetStateLabel(stateIndex);

                domain.Propagate(Domain.Equilibrium.H_EQUILIBRIUM_SUM, Domain.EvidenceMode.H_EVIDENCE_MODE_NORMAL);

                if (i != GetFirstTimeStep(domain))
                {
                    DiscreteChanceNode observationNode = (DiscreteChanceNode)domain.GetNodeByName(ojPrefix + i);
                    uint obsIndex = observationNode.SetRandomState();
                    result.observation = observationNode.GetStateLabel(obsIndex);
                }
                domain.RetractStateHistories(sPrefix);
                domain.Propagate(Domain.Equilibrium.H_EQUILIBRIUM_SUM, Domain.EvidenceMode.H_EVIDENCE_MODE_NORMAL);

                DiscreteDecisionNode decisionNode = (DiscreteDecisionNode) domain.GetNodeByName(ajPrefix + i);
                uint decIndex = decisionNode.SetStateWithHighestReward();
                result.action = decisionNode.GetStateLabel(decIndex);

                domain.ResetStateHistories(stateHistory, sPrefix);
                domain.Propagate(Domain.Equilibrium.H_EQUILIBRIUM_SUM, Domain.EvidenceMode.H_EVIDENCE_MODE_NORMAL);

                UtilityNode utilityNode = (UtilityNode)domain.GetNodeByName(uPrefix + i);
                result.util = utilityNode.GetExpectedUtility();
                
                results.Add(result);
            }

            domain.RetractFindings();
            return results;
        }

        /// <summary>
        /// Simulates an I-DID domain by including simulating a DID domain
        /// Simultaneously
        /// </summary>
        /// <param name="domain">domain</param>
        /// <param name="jDomain">agent j domain</param>
        /// <param name="ajPrefix">agent j action prefix</param>
        /// <param name="ojPrefix">agent j observation prefix</param>
        /// <param name="aiPrefix">agent i action prefix</param>
        /// <param name="oiPrefix">agent i observation prefix</param>
        /// <param name="sPrefix">state prefix</param>
        /// <param name="uPrefix">utility</param>
        /// <param name="model">model</param>
        /// <returns>results</returns>
        public static List<SimulationResult> GetSimulationResultIdid(this Domain domain, Domain jDomain, string ajPrefix, string ojPrefix,
            string aiPrefix, string oiPrefix, string sPrefix, string uPrefix, string model)
        {
            List<SimulationResult> results = new List<SimulationResult>();

            List<uint> stateHistory = new List<uint>();
            List<uint> ajHistory = new List<uint>();

            for (int i = GetFirstTimeStep(domain); i <= GetLastTimeStep(domain); i++)
            {
                SimulationResult result = new SimulationResult();
                result.model = model;
                result.timeStep = i;

                try
                {
                    DiscreteChanceNode stateNodeI = (DiscreteChanceNode)domain.GetNodeByName(sPrefix + i);
                    DiscreteChanceNode stateNodeJ = (DiscreteChanceNode)jDomain.GetNodeByName(sPrefix + i);
                    uint stateIndex = stateNodeI.SetRandomState();
                    stateNodeJ.SelectState(stateIndex);
                    stateHistory.Add(stateIndex);
                    result.state = stateNodeI.GetStateLabel(stateIndex);

                    domain.Propagate(Domain.Equilibrium.H_EQUILIBRIUM_SUM, Domain.EvidenceMode.H_EVIDENCE_MODE_NORMAL);
                    jDomain.Propagate(Domain.Equilibrium.H_EQUILIBRIUM_SUM, Domain.EvidenceMode.H_EVIDENCE_MODE_NORMAL);

                    if (i != GetFirstTimeStep(domain))
                    {
                        DiscreteChanceNode observationNodeI = (DiscreteChanceNode)domain.GetNodeByName(oiPrefix + i);
                        DiscreteChanceNode observationNodeJ = (DiscreteChanceNode)domain.GetNodeByName(ojPrefix + i);
                        observationNodeJ.SetRandomState();
                        uint obsIndexI = observationNodeI.SetRandomState();
                        result.observation = observationNodeI.GetStateLabel(obsIndexI);
                    }
                    domain.RetractStateHistories(sPrefix);
                    domain.RetractStateHistories(ajPrefix);
                    jDomain.RetractStateHistories(sPrefix);

                    domain.Propagate(Domain.Equilibrium.H_EQUILIBRIUM_SUM, Domain.EvidenceMode.H_EVIDENCE_MODE_NORMAL);
                    jDomain.Propagate(Domain.Equilibrium.H_EQUILIBRIUM_SUM, Domain.EvidenceMode.H_EVIDENCE_MODE_NORMAL);

                    DiscreteDecisionNode decisionNodeI = (DiscreteDecisionNode)domain.GetNodeByName(aiPrefix + i);
                    DiscreteDecisionNode decisionNodeJ = (DiscreteDecisionNode)jDomain.GetNodeByName(ajPrefix + i);

                    uint decIndexI = decisionNodeI.SetStateWithHighestReward();
                    uint decindexJ = decisionNodeJ.SetStateWithHighestReward();

                    ajHistory.Add(decindexJ);

                    result.action = decisionNodeI.GetStateLabel(decIndexI);

                    jDomain.ResetStateHistories(stateHistory, sPrefix);
                    domain.ResetStateHistories(stateHistory, sPrefix);
                    //domain.ResetStateHistories(ajHistory, ajPrefix);

                    domain.Propagate(Domain.Equilibrium.H_EQUILIBRIUM_SUM, Domain.EvidenceMode.H_EVIDENCE_MODE_NORMAL);

                    UtilityNode utilityNode = (UtilityNode)domain.GetNodeByName(uPrefix + i);
                    result.util = utilityNode.GetExpectedUtility();
                }
                catch (ExceptionHugin eh)
                {
                    result.util = -200;
                    //return new List<SimulationResult>();
                }

                results.Add(result);
            }


            domain.RetractFindings();
            jDomain.RetractFindings();
            return results;
        }

        /// <summary>
        /// retracts all state evidence
        /// </summary>
        /// <param name="domain">domain</param>
        /// <param name="sPrefix">state prefix</param>
        public static void RetractStateHistories(this Domain domain, string sPrefix)
        {
            for (int i = GetFirstTimeStep(domain); i <= GetLastTimeStep(domain); i++)
            {
                DiscreteChanceNode stateNode = (DiscreteChanceNode)domain.GetNodeByName(sPrefix + i);
                stateNode.RetractFindings();
            }
        }

        /// <summary>
        /// Resets the state histories for calculating next state and rewards
        /// </summary>
        /// <param name="domain">domain</param>
        /// <param name="history">history to reset to</param>
        /// <param name="sPrefix">state prefix</param>
        public static void ResetStateHistories(this Domain domain, List<uint> history, string sPrefix)
        {
            for (int i = 0; i < history.Count; i++)
            {
                DiscreteChanceNode stateNode = (DiscreteChanceNode)domain.GetNodeByName(sPrefix + (i + domain.GetFirstTimeStep()));
                stateNode.SelectState(history[i]);
            }
        }
    }
}
